#!/usr/bin/env python3
"""
数据库迁移脚本 - 添加AI总结字段
为EmailMessage表添加ai_summary字段
"""

import sqlite3
import os
from pathlib import Path

def migrate_database():
    """执行数据库迁移"""
    
    # 数据库文件路径
    db_path = Path(__file__).parent.parent / "instance" / "feng_admin.db"
    
    if not db_path.exists():
        print(f"数据库文件不存在: {db_path}")
        print("请先运行后端服务以创建数据库")
        return False
    
    try:
        # 连接数据库
        conn = sqlite3.connect(str(db_path))
        cursor = conn.cursor()
        
        # 检查ai_summary字段是否已存在
        cursor.execute("PRAGMA table_info(email_messages)")
        columns = [column[1] for column in cursor.fetchall()]
        
        if 'ai_summary' in columns:
            print("ai_summary字段已存在，无需迁移")
            return True
        
        # 添加ai_summary字段
        print("正在添加ai_summary字段...")
        cursor.execute("ALTER TABLE email_messages ADD COLUMN ai_summary TEXT")
        
        # 提交更改
        conn.commit()
        print("✅ ai_summary字段添加成功")
        
        # 验证字段是否添加成功
        cursor.execute("PRAGMA table_info(email_messages)")
        columns = [column[1] for column in cursor.fetchall()]
        
        if 'ai_summary' in columns:
            print("✅ 迁移完成，ai_summary字段已成功添加")
            return True
        else:
            print("❌ 迁移失败，ai_summary字段未找到")
            return False
            
    except sqlite3.Error as e:
        print(f"❌ 数据库操作失败: {e}")
        return False
    except Exception as e:
        print(f"❌ 迁移过程中发生错误: {e}")
        return False
    finally:
        if 'conn' in locals():
            conn.close()

def rollback_migration():
    """回滚迁移（删除ai_summary字段）"""
    
    db_path = Path(__file__).parent.parent / "instance" / "feng_admin.db"
    
    if not db_path.exists():
        print(f"数据库文件不存在: {db_path}")
        return False
    
    try:
        conn = sqlite3.connect(str(db_path))
        cursor = conn.cursor()
        
        # 检查ai_summary字段是否存在
        cursor.execute("PRAGMA table_info(email_messages)")
        columns = [column[1] for column in cursor.fetchall()]
        
        if 'ai_summary' not in columns:
            print("ai_summary字段不存在，无需回滚")
            return True
        
        # SQLite不支持DROP COLUMN，需要重建表
        print("⚠️  警告: SQLite不支持直接删除列，需要重建表")
        print("这将删除所有AI总结数据，是否继续？(y/N): ", end="")
        
        confirm = input().strip().lower()
        if confirm != 'y':
            print("回滚已取消")
            return False
        
        # 创建新表（不包含ai_summary字段）
        cursor.execute("""
            CREATE TABLE email_messages_new (
                id INTEGER PRIMARY KEY,
                user_id INTEGER NOT NULL,
                email_account_id VARCHAR(100) NOT NULL,
                message_id VARCHAR(200) NOT NULL,
                imap_uid INTEGER,
                subject VARCHAR(500),
                sender VARCHAR(200),
                recipients TEXT,
                cc TEXT,
                body_text TEXT,
                body_html TEXT,
                attachments TEXT,
                is_read BOOLEAN DEFAULT 0,
                is_starred BOOLEAN DEFAULT 0,
                is_full_content BOOLEAN DEFAULT 0,
                content_size INTEGER,
                folder VARCHAR(50) DEFAULT 'inbox',
                imap_folder VARCHAR(200),
                received_at DATETIME,
                created_at DATETIME DEFAULT CURRENT_TIMESTAMP,
                updated_at DATETIME DEFAULT CURRENT_TIMESTAMP,
                FOREIGN KEY (user_id) REFERENCES users (id),
                FOREIGN KEY (email_account_id) REFERENCES email_accounts (id)
            )
        """)
        
        # 复制数据（排除ai_summary字段）
        cursor.execute("""
            INSERT INTO email_messages_new 
            SELECT id, user_id, email_account_id, message_id, imap_uid, subject, sender, 
                   recipients, cc, body_text, body_html, attachments, is_read, is_starred, 
                   is_full_content, content_size, folder, imap_folder, received_at, 
                   created_at, updated_at
            FROM email_messages
        """)
        
        # 删除旧表
        cursor.execute("DROP TABLE email_messages")
        
        # 重命名新表
        cursor.execute("ALTER TABLE email_messages_new RENAME TO email_messages")
        
        # 重新创建索引
        cursor.execute("CREATE INDEX idx_message_id ON email_messages (message_id)")
        cursor.execute("CREATE INDEX idx_email_account_id ON email_messages (email_account_id)")
        cursor.execute("CREATE INDEX idx_folder ON email_messages (folder)")
        cursor.execute("CREATE INDEX idx_is_read ON email_messages (is_read)")
        cursor.execute("CREATE INDEX idx_imap_uid ON email_messages (imap_uid)")
        cursor.execute("CREATE INDEX idx_is_full_content ON email_messages (is_full_content)")
        cursor.execute("CREATE INDEX idx_user_folder_read ON email_messages (user_id, folder, is_read)")
        cursor.execute("CREATE INDEX idx_user_received_at ON email_messages (user_id, received_at)")
        cursor.execute("CREATE INDEX idx_user_account_folder ON email_messages (user_id, email_account_id, folder)")
        cursor.execute("CREATE UNIQUE INDEX unique_user_message ON email_messages (user_id, message_id)")
        
        conn.commit()
        print("✅ 回滚完成，ai_summary字段已删除")
        return True
        
    except sqlite3.Error as e:
        print(f"❌ 数据库操作失败: {e}")
        return False
    except Exception as e:
        print(f"❌ 回滚过程中发生错误: {e}")
        return False
    finally:
        if 'conn' in locals():
            conn.close()

if __name__ == "__main__":
    import sys
    
    if len(sys.argv) > 1 and sys.argv[1] == "rollback":
        print("🔄 开始回滚迁移...")
        success = rollback_migration()
    else:
        print("🚀 开始数据库迁移...")
        success = migrate_database()
    
    if success:
        print("✅ 操作完成")
        sys.exit(0)
    else:
        print("❌ 操作失败")
        sys.exit(1)
